import sqlite3
import uuid
import image_pro.config as config

class ImageDB:
    def __init__(self, db_name=config.DB_NAME):
        self.db_name = db_name
        self.conn = self.create_connection()
        self.init_db()
        
    def create_connection(self):
        """创建数据库连接"""
        return sqlite3.connect(self.db_name, check_same_thread=False)
    
    def close_connection(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
    
    def execute(self, sql, params=()):
        """执行SQL语句"""
        cursor = self.conn.cursor()
        try:
            cursor.execute(sql, params)
            self.conn.commit()
            return cursor
        except sqlite3.Error as e:
            self.conn.rollback()
            raise e
    
    def generate_uid(self):
        """生成唯一ID"""
        return str(uuid.uuid4())
    
    def init_db(self):
        """初始化数据库表结构"""
        sqls = [
            f'''
            CREATE TABLE IF NOT EXISTS {config.TABLE_TASKS} (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                refer TEXT,
                path TEXT DEFAULT './',
                url TEXT,
                source_type TEXT NOT NULL,
                uid TEXT NOT NULL UNIQUE,
                state TEXT DEFAULT '{config.STATUS_PENDING}',
                total_images INTEGER DEFAULT 0,
                downloaded INTEGER DEFAULT 0,
                create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
            ''',
            f'''
            CREATE TABLE IF NOT EXISTS {config.TABLE_DOWNLOADS} (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_uid TEXT NOT NULL,
                url TEXT NOT NULL,
                path TEXT NOT NULL,
                filename TEXT NOT NULL,
                state TEXT DEFAULT '{config.STATUS_PENDING}',
                create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY(task_uid) REFERENCES {config.TABLE_TASKS}(uid)
            )
            '''
        ]
        
        for sql in sqls:
            self.execute(sql)
    
    def create_image_task(self, **kwargs):
        """创建图片下载任务"""
        uid = self.generate_uid()
        sql = f'''
        INSERT INTO {config.TABLE_TASKS} 
        (name, refer, path, url, source_type, uid, total_images, state) 
        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        '''
        params = (
            kwargs['name'],
            kwargs.get('refer', ''),
            kwargs['path'],
            kwargs.get('url', ''),
            kwargs['source_type'],
            uid,
            kwargs.get('total_images', 0),
            config.STATUS_READY
        )
        cursor = self.execute(sql, params)
        return uid, cursor.lastrowid
    
    def create_image_download(self, **kwargs):
        """创建图片下载项"""
        sql = f'''
        INSERT INTO {config.TABLE_DOWNLOADS} 
        (task_uid, url, path, filename, state) 
        VALUES (?, ?, ?, ?, ?)
        '''
        params = (
            kwargs['task_uid'],
            kwargs['url'],
            kwargs['path'],
            kwargs['filename'],
            config.STATUS_PENDING
        )
        cursor = self.execute(sql, params)
        return cursor.lastrowid
    
    def get_all_tasks(self):
        """获取所有任务"""
        sql = f'''
        SELECT id, uid, name, state, total_images, downloaded, path
        FROM {config.TABLE_TASKS}
        ORDER BY create_time DESC
        '''
        cursor = self.execute(sql)
        return cursor.fetchall()
    
    def get_task_by_uid(self, uid):
        """通过UID获取任务详情"""
        sql = f'''
        SELECT * FROM {config.TABLE_TASKS}
        WHERE uid = ?
        '''
        cursor = self.execute(sql, (uid,))
        return cursor.fetchone()
    
    def get_task_downloads(self, task_uid):
        """获取任务的所有下载项"""
        sql = f'''
        SELECT * FROM {config.TABLE_DOWNLOADS}
        WHERE task_uid = ?
        ORDER BY id ASC
        '''
        cursor = self.execute(sql, (task_uid,))
        return cursor.fetchall()
    
    def get_pending_downloads(self, task_uid, limit=10):
        """获取待下载的项目"""
        sql = f'''
        SELECT * FROM {config.TABLE_DOWNLOADS}
        WHERE task_uid = ? AND state = ?
        ORDER BY id ASC
        LIMIT ?
        '''
        cursor = self.execute(sql, (task_uid, config.STATUS_PENDING, limit))
        return cursor.fetchall()
    
    def update_download_state(self, download_id, state):
        """更新下载项状态"""
        sql = f'''
        UPDATE {config.TABLE_DOWNLOADS}
        SET state = ?, update_time = CURRENT_TIMESTAMP
        WHERE id = ?
        '''
        self.execute(sql, (state, download_id))
    
    def update_task_state(self, task_uid, state):
        """更新任务状态"""
        sql = f'''
        UPDATE {config.TABLE_TASKS}
        SET state = ?, update_time = CURRENT_TIMESTAMP
        WHERE uid = ?
        '''
        self.execute(sql, (state, task_uid))
    
    def update_task_progress(self, task_uid, downloaded):
        """更新任务下载进度"""
        sql = f'''
        UPDATE {config.TABLE_TASKS}
        SET downloaded = ?, update_time = CURRENT_TIMESTAMP
        WHERE uid = ?
        '''
        self.execute(sql, (downloaded, task_uid))
    
    def delete_task(self, task_uid):
        """删除任务及其下载项"""
        # 先删除下载项
        sql1 = f'''
        DELETE FROM {config.TABLE_DOWNLOADS}
        WHERE task_uid = ?
        '''
        self.execute(sql1, (task_uid,))
        
        # 再删除任务
        sql2 = f'''
        DELETE FROM {config.TABLE_TASKS}
        WHERE uid = ?
        '''
        self.execute(sql2, (task_uid,))